// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/process/stack_trace.h>

// Symbolizer requires melon to be compiled with libelf and libdwarf support
// (also currently only works in Linux).
#if __linux__ && MELON_HAVE_ELF && MELON_HAVE_DWARF
#define POLLUX_HAS_SYMBOLIZER 1
#else
#define POLLUX_HAS_SYMBOLIZER 0
#endif

#include <algorithm>
#include <fstream>

#include <fmt/format.h>
#include <melon/indestructible.h>
#include <melon/string.h>
#include <melon/experimental/symbolizer/stack_trace.h>

#include <pollux/common/process/process_base.h>

#ifdef __linux__
#include <melon/experimental/symbolizer/symbolizer.h> // @manual
#include <melon/fibers/fiber_manager.h> // @manual
#endif

namespace kumo::pollux::process {
    StackTrace::StackTrace(int32_t skipFrames) {
        create(skipFrames);
    }

    StackTrace::StackTrace(const StackTrace &other) {
        bt_pointers_ = other.bt_pointers_;
        if (melon::test_once(other.bt_vector_flag_)) {
            bt_vector_ = other.bt_vector_;
            melon::call_once(bt_vector_flag_, [] {
            }); // Set the flag.
        }
        if (melon::test_once(other.bt_flag_)) {
            bt_ = other.bt_;
            melon::call_once(bt_flag_, [] {
            }); // Set the flag.
        }
    }

    StackTrace &StackTrace::operator=(const StackTrace &other) {
        if (this != &other) {
            this->~StackTrace();
            new(this) StackTrace(other);
        }
        return *this;
    }

    void StackTrace::create(int32_t skipFrames) {
        const int32_t kDefaultSkipFrameAdjust = 2; // ::create(), ::StackTrace()
        const int32_t kMaxFrames = 75;

        bt_pointers_.clear();
        uintptr_t btpointers[kMaxFrames];
        ssize_t framecount = melon::symbolizer::getStackTrace(btpointers, kMaxFrames);
        if (framecount <= 0) {
            return;
        }

        framecount = std::min(framecount, static_cast<ssize_t>(kMaxFrames));
        skipFrames = std::max(skipFrames + kDefaultSkipFrameAdjust, 0);

        bt_pointers_.reserve(framecount - skipFrames);
        for (int32_t i = skipFrames; i < framecount; i++) {
            bt_pointers_.push_back(reinterpret_cast<void *>(btpointers[i]));
        }
    }

    ///////////////////////////////////////////////////////////////////////////////
    // reporting functions

    const std::vector<std::string> &StackTrace::toStrVector() const {
        melon::call_once(bt_vector_flag_, [&] {
            size_t frame = 0;
            static melon::Indestructible<melon::kmstring> myname{
                melon::demangle(typeid(decltype(*this))) + "::"
            };
            bt_vector_.reserve(bt_pointers_.size());
            for (auto ptr: bt_pointers_) {
                auto framename = translateFrame(ptr);
                if (melon::StringPiece(framename).startsWith(*myname)) {
                    continue; // ignore frames in the StackTrace class
                }
                bt_vector_.push_back(fmt::format("# {:<2d} {}", frame++, framename));
            }
        });
        return bt_vector_;
    }

    const std::string &StackTrace::toString() const {
        melon::call_once(bt_flag_, [&] {
            const auto &vec = toStrVector();
            size_t needed = 0;
            for (const auto &frame: vec) {
                needed += frame.size() + 1;
            }
            bt_.reserve(needed);
            for (const auto &frame_title: vec) {
                bt_ += frame_title;
                bt_ += '\n';
            }
        });
        return bt_;
    }

    std::string StackTrace::log(
        const char *errorType,
        std::string *out /* = NULL */) const {
        std::string pid = melon::to<std::string>(getProcessId());

        std::string msg;
        msg += "Host: " + getHostName();
        msg += "\nProcessID: " + pid;
        msg += "\nThreadID: " +
                melon::to<std::string>(reinterpret_cast<uintptr_t>(getThreadId()));
        msg += "\nName: " + getAppName();
        msg += "\nType: ";
        if (errorType) {
            msg += errorType;
        } else {
            msg += "(unknown error)";
        }
        msg += "\n\n";
        msg += toString();
        msg += "\n";

        std::string tracefn = "/tmp/stacktrace." + pid + ".log";
        std::ofstream f(tracefn.c_str());
        if (f) {
            f << msg;
            f.close();
        }

        if (out) {
            *out = msg;
        }
        return tracefn;
    }

#if POLLUX_HAS_SYMBOLIZER
namespace {
inline std::string translateFrameImpl(void* addressPtr) {
  // TODO: lineNumbers has been disabled since 2009.
  using namespace melon::symbolizer;

  std::uintptr_t address = reinterpret_cast<std::uintptr_t>(addressPtr);
  Symbolizer symbolizer(LocationInfoMode::DISABLED);
  SymbolizedFrame frame;
  symbolizer.symbolize(address, frame);

  StringSymbolizePrinter printer(SymbolizePrinter::TERSE);
  printer.print(frame);
  return printer.str();
}
} // namespace
#endif

    std::string StackTrace::translateFrame(void *addressPtr, bool /*lineNumbers*/) {
#if POLLUX_HAS_SYMBOLIZER
  return melon::fibers::runInMainContext(
      [addressPtr]() { return translateFrameImpl(addressPtr); });
#else
        static_cast<void>(addressPtr);
        return std::string{};
#endif
    }

    std::string StackTrace::demangle(const char *mangled) {
        return melon::demangle(mangled).toStdString();
    }
} // namespace kumo::pollux::process
